import torchvision
import torch
import torchvision.transforms as transforms
import os, glob
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import cv2
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

preprocessing_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225]),
])

preprocessing_val = transforms.Compose([    #[1]
 transforms.Resize(256),                    #[2]
 transforms.CenterCrop(224),                #[3]
 transforms.ToTensor(),                     #[4]
 transforms.Normalize(                      #[5]
 mean=[0.485, 0.456, 0.406],                #[6]
 std=[0.229, 0.224, 0.225]                  #[7]
 )])

model = torchvision.models.resnet34(pretrained=True)
model.eval()
a = os.listdir("/media/retoo/RetooDisk/wanghui/Data/ILSVR2012C/ILSVRC2012_img_val/")
category = open("../imagenet1000_clsidx_to_labels.txt", "r").readlines()
id_to_label = dict()
for line in category:
  line = line.split(":")
  id_to_label[int(line[0])] = line[1].strip()

root = "/media/retoo/RetooDisk/wanghui/Data/ILSVR2012C/ILSVRC2012_img_val/"
files = open("/media/retoo/RetooDisk/wanghui/Data/ILSVR2012C/ILSVRC2012_val.txt","r")
for line in files.readlines():
  line = line.strip().split(" ")
  gt_label = int(line[1])
  p = os.path.join(root,line[0].split("/")[-1])

  #official preprocess
  # src = Image.open(p).convert('RGB')
  # img = preprocessing_val(src)

  #modify proprocess
  src = cv2.imread(p)
  src = cv2.cvtColor(src,cv2.COLOR_BGR2RGB)
  src = cv2.resize(src, (224,224))
  img = torch.from_numpy(np.asarray(src)).permute((2, 0, 1)).contiguous()
  img = img/255.

  out = model(img.unsqueeze(dim=0))
  label = torch.argmax(out, dim=-1).numpy()
  plt.imshow(src)
  plt.show()
  print("pred: ", id_to_label[label[0]], " true: ", id_to_label[gt_label])
